package org.seed.rpc.core.etcd;

import com.google.common.util.concurrent.ListenableFuture;
import com.ibm.etcd.api.*;
import com.ibm.etcd.client.KvStoreClient;
import com.ibm.etcd.client.ListenerObserver;
import com.ibm.etcd.client.kv.KvClient;
import com.ibm.etcd.client.kv.WatchUpdate;
import com.ibm.etcd.client.lease.LeaseClient;
import com.ibm.etcd.client.lock.LockClient;
import io.grpc.stub.StreamObserver;
import lombok.extern.slf4j.Slf4j;
import org.seed.rpc.core.RegistryClient;
import com.google.protobuf.ByteString;
import org.seed.rpc.core.Constants;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;


@Slf4j
public class EtcdClient implements RegistryClient {

    //初始化连接
    private KvStoreClient kvStoreClient;

    //对数据操作的连接
    private KvClient kvClient;

    //对租期等时间操作的连接
    private LeaseClient leaseClient;

    //对分布式锁操作的连接
    private LockClient lockClient;

    /**
     * 在Servlet容器初始化前执行
     */
    public EtcdClient(String endpoint) {
        kvStoreClient = com.ibm.etcd.client.EtcdClient.forEndpoints(endpoint).withPlainText().build();
        this.kvClient = kvStoreClient.getKvClient();
        this.leaseClient = kvStoreClient.getLeaseClient();
        this.lockClient = kvStoreClient.getLockClient();
    }

    /**
     * 添加操作，永不过期
     * 形式: 同步
     *
     * @param key
     * @param value
     */
    @Override
    public void put(String key, String value) {
        kvClient.put(ByteString.copyFromUtf8(key), ByteString.copyFromUtf8(value)).sync();
    }

    /**
     * 添加操作，设置租期ID
     *
     * @param key
     * @param value
     * @param leaseId
     */
    private void put(String key, String value, long leaseId) {
        kvClient.put(ByteString.copyFromUtf8(key), ByteString.copyFromUtf8(value), leaseId).sync();
    }

    /**
     * 存入key、value，和过期时间，单位是秒
     */
    public Long putAndGrant(String key, String value, long ttl) {
        LeaseGrantResponse lease = leaseClient.grant(ttl).sync();
        long leaseId = lease.getID();
        put(key, value, leaseId);
        return leaseId;
    }


    /**
     * 执行续租
     *
     * @param leaseId 租约ID
     */
    public void keepAliveOnce(long leaseId) {
        leaseClient.keepAliveOnce(leaseId);
    }


    /**
     * 删除租用
     *
     * @param leaseId 租约ID
     */
    public void revoke(long leaseId) {
        leaseClient.revoke(leaseId);
    }


    /**
     * 根据key，获取value
     */
    @Override
    public String get(String key) {
        RangeResponse rangeResponse = kvClient.get(ByteString.copyFromUtf8(key)).sync();
        List<KeyValue> keyValues = rangeResponse.getKvsList();
        if (keyValues == null || keyValues.size() == 0) {
            return null;
        }
        return keyValues.get(0).getValue().toStringUtf8();
    }

    /**
     * 查询指定前缀的数据
     *
     * @param prefix
     */
    @Override
    public Map<String, String> getPrefix(String prefix) {
        try {
            RangeRequest rangeRequest = RangeRequest.newBuilder()
                    .setRangeEnd(ByteString.copyFromUtf8("\0"))
                    .setKey(ByteString.copyFromUtf8(prefix))
                    .build();
            ListenableFuture<RangeResponse> future = kvClient.get(rangeRequest);
            // 范围查询的取值
            RangeResponse response = future.get();
            Map<String, String> rst = new HashMap<>();
            List<KeyValue> kvList = response.getKvsList();
            if (response == null || kvList == null || kvList.size() == 0) {
                return rst;
            }
            for (KeyValue kv : kvList) {
                rst.put(kv.getKey().toStringUtf8(), kv.getValue().toStringUtf8());
            }
            return rst;
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        }
    }

    /**
     * 删除指定key的数据
     */
    @Override
    public void remove(String key) {
        kvClient.delete(ByteString.copyFromUtf8(key)).sync();
    }

    /**
     * 监听值的变化
     *
     * @param prefix   键
     * @param consumer 内容变化后执行的方法
     */
    @Override
    public void prefixOnChange(final String prefix, Consumer consumer) {
        StreamObserver<WatchUpdate> observer = new ListenerObserver<WatchUpdate>() {
            @Override
            public void event(boolean b, WatchUpdate watchUpdate, Throwable throwable) {
                List<Event> events = watchUpdate.getEvents();
                if (events == null || events.size() == 0) {
                    return;
                }
                boolean updateAction = false;
                String eventName = null;
                for (Event event : events) {
                    if (event.getType() == Event.EventType.DELETE || event.getType() == Event.EventType.PUT || event.getType() == Event.EventType.UNRECOGNIZED) {
                        updateAction = true;
                        eventName = event.getType().name();
                        break;
                    }
                }
                if (updateAction) {
                    log.info("received etcd event: {}, refresh local config", eventName);
                    consumer.accept(prefix);
                }
            }
        };
        WatchCreateRequest request = WatchCreateRequest.newBuilder()
                .setRangeEnd(ByteString.copyFromUtf8("\0"))
                .setKey(ByteString.copyFromUtf8(Constants.SERVICE_REGISTRY_KEY_PREFIX))
                .build();
        KvClient.Watch watch = kvClient.watch(request, observer);
        log.info("client watch, is cancelled: {},is done: {}", watch.isCancelled(), watch.isDone());
    }


}
